package org.jeecg.modules.openai.service;

import cn.hutool.http.HttpUtil;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONObject;
import com.theokanning.openai.completion.CompletionChoice;
import com.theokanning.openai.completion.CompletionRequest;
import com.theokanning.openai.completion.CompletionResult;
import com.theokanning.openai.image.CreateImageEditRequest;
import com.theokanning.openai.image.CreateImageRequest;
import com.theokanning.openai.image.ImageResult;
import com.theokanning.openai.service.OpenAiService;
import lombok.extern.slf4j.Slf4j;
import org.apache.shiro.SecurityUtils;
import org.jeecg.common.exception.JeecgBootException;
import org.jeecg.common.system.vo.LoginUser;
import org.jeecg.modules.openai.entity.*;
import org.jeecg.modules.openai.entity.AiImgCreateRecord;
import org.jeecg.modules.openai.model.ChatModel;
import org.jeecg.modules.openai.model.EditImg;
import org.jeecg.modules.openai.model.RedBook;
import org.jeecg.modules.openai.model.WriteModel;
import org.jeecg.modules.openai.util.FileUtils;
import org.jeecg.modules.openai.util.GPT2Tokenizer;
import org.jeecg.modules.system.entity.SysUser;
import org.jeecg.modules.system.service.ISysUserService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.ClassPathResource;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.math.BigDecimal;
import java.net.HttpURLConnection;
import java.net.MalformedURLException;
import java.net.URL;
import java.nio.file.Paths;
import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

@Service
@Slf4j
public class OpenApiService {

    private static String organization = "org-tEbuhtKZqfRM9Q3uXzddiATb";
    private static String apiKey = "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";

    private static int avgToken = 1000 / 750;

    private static Map<String, BigDecimal> modelPrice = new HashMap<>();

    static {

        modelPrice.put("text-davinci-003", new BigDecimal("0.14"));
        modelPrice.put("gpt-3.5-turbo", new BigDecimal("0.014"));
        modelPrice.put("gpt-3.5-turbo-0301", new BigDecimal("0.014"));

    }

    @Autowired
    private ISysUserService sysUserService;

    public List<String> chat(String input) {

        List<String> stop = new ArrayList<>();

        stop.add("Human:");

        stop.add("AI:");

        CompletionResult result = this.completion(input, stop, getMaxTokens("text-davinci-003", input), "text-davinci-003");

        this.handelResultToken(result);

        return result.getChoices().stream().map(i -> i.getText()).collect(Collectors.toList());
    }



    public String write(String prompt) {


        List<GPTMessage> list = new ArrayList<>();

        GPTMessage message = new GPTMessage();

        message.setRole("user");
        message.setContent(prompt);
        list.add(message);
        GPTCompletionRequest request = new GPTCompletionRequest();
        request.setModel("gpt-3.5-turbo");
        request.setMessages(list);
        request.setMax_tokens(getMaxTokens("gpt-3.5-turbo", prompt));
        String rp = HttpUtil.createPost("https://pms.descartes.digital/openai/v1/chat/completions")
                .timeout(1000*60).body(new JSONObject(request).toString(), "application/json")
                .bearerAuth(apiKey).execute().body();
        GPTCompletionResult result = new JSONObject(rp).toBean(GPTCompletionResult.class);
        handelResultToken(result);
//        cn.hutool.core.bean.copier.CopyOptions x =
        return ((GPTCompletionChoice)result.getChoices().get(0)).getMessage().getContent();
    }

    public String redBook(RedBook redBook){

        List<GPTMessage> list = new ArrayList<>();
        GPTMessage message = new GPTMessage();
        message.setRole("user");
        message.setContent("请根据我下面的提示帮我生成文案");
        list.add(message);
        message = new GPTMessage();
        message.setRole("user");
        message.setContent("写一篇关于\""+redBook.getPrompt()+"\"的文案");
        list.add(message);
        message = new GPTMessage();
        message.setRole("user");
        message.setContent("请你以小红书博主的文案写作方式，用词特点为参考，在文案中合适的地方尽可能多的加入emoji表情并且在文案结尾加入这个文案相关话题引用");
        list.add(message);
        message = new GPTMessage();
        message.setRole("user");
        message.setContent("在文案中加入类似于"+redBook.getKeys()+"的内容，将它们巧妙地加入到文章内容中");
        list.add(message);

        GPTCompletionRequest request = new GPTCompletionRequest();
        request.setModel("gpt-3.5-turbo");
        request.setMessages(list);
        request.setMax_tokens(getMaxTokens("gpt-3.5-turbo", redBook.getPrompt()));

        String rp = HttpUtil.createPost("https://pms.descartes.digital/openai/v1/chat/completions")
                .timeout(1000*60).body(new JSONObject(request).toString(), "application/json")
                .bearerAuth(apiKey).execute().body();
        GPTCompletionResult result = new JSONObject(rp).toBean(GPTCompletionResult.class);
        handelResultToken(result);
//        cn.hutool.core.bean.copier.CopyOptions x =
        return ((GPTCompletionChoice)result.getChoices().get(0)).getMessage().getContent();


    }

    public List<String> gptChat(String input) {


        String regx = "(我:.*)|(AI:.*)";



        Pattern pat = Pattern.compile(regx);
        Matcher mat = pat.matcher(input);
        List<GPTMessage> list = new ArrayList<>();

        while (mat.find()) {
            String[] str = mat.group().split(":+?");
            GPTMessage message = new GPTMessage();
            message.setRole("我".equals(str[0])?"user":"assistant");
            String content = "";
            for (int i = 1; i < str.length; i++) {
                content += str[i];
            }
            message.setContent(content);
            list.add(message);
        }

        GPTCompletionRequest request = new GPTCompletionRequest();
        request.setModel("gpt-3.5-turbo");
        request.setMessages(list);
        request.setMax_tokens(getMaxTokens("gpt-3.5-turbo", input));
        String rp = HttpUtil.createPost("https://pms.descartes.digital/openai/v1/chat/completions")
                .timeout(1000*60).body(new JSONObject(request).toString(), "application/json")
                .bearerAuth(apiKey).execute().body();
        GPTCompletionResult result = new JSONObject(rp).toBean(GPTCompletionResult.class);
        handelResultToken(result);
//        cn.hutool.core.bean.copier.CopyOptions x =
        return result.getChoices().stream().map(i -> {
            GPTCompletionChoice c = (GPTCompletionChoice) i;
            return c.getMessage().getContent();
        }).collect(Collectors.toList());
    }


    public Integer getMaxTokens(String model, String prompt) {

        LoginUser sysUser = (LoginUser) SecurityUtils.getSubject().getPrincipal();

        SysUser user = sysUserService.getById(sysUser.getId());

        int score = user.getScore(); //用户余额

        BigDecimal price = modelPrice.get(model);

        System.out.println(new ClassPathResource(Paths.get("/tokenizers/gpt2", "encoder.json").toString()));
        GPT2Tokenizer tokenizer = null;
        try {
            tokenizer = GPT2Tokenizer.fromPretrained("/tokenizers/gpt2");
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        List<Integer> result = tokenizer.encode(prompt);
        //提示需要消耗的token
        int promptTokens = result.size();

        log.info("用户{},提示token:{}", sysUser.getId(), promptTokens);

        //计算用户剩余token
        int userToken = new BigDecimal(score).divide(price, BigDecimal.ROUND_HALF_UP).intValue();

        if (userToken - promptTokens <= 0) {
            throw new JeecgBootException("积分不足");
        }
        return userToken > 2000 ? 2000 : userToken;
    }

    public void handelResultToken(CompletionResult result) {


        long totalTokens = result.getUsage().getTotalTokens();

        LoginUser sysUser = (LoginUser) SecurityUtils.getSubject().getPrincipal();

        log.info("usage:id:{},data:{}", sysUser.getId(), result.getUsage());

        SysUser user = sysUserService.getById(sysUser.getId());

        BigDecimal price = modelPrice.get(result.getModel());

        int score = new BigDecimal(totalTokens).multiply(price).intValue(); //用户余额

        if (user.getScore().compareTo(score) >= 0) {
            user.setScore(user.getScore() - score);
        } else {
            user.setScore(0);
        }

        sysUserService.updateById(user);

    }

    public List<String> naturalLanguageToSQL(String input) {

        Duration timeout = Duration.ofSeconds(60);
        OpenAiService service = new OpenAiService(apiKey, timeout);
        Integer maxTokens = 800;
        List<String> stop = new ArrayList<>();
        stop.add("#");
        stop.add(";");
        CompletionRequest completionRequest = CompletionRequest.builder()
                .prompt(input)
                .user("user123")
                .stop(stop)
                .topP(1.0)
                .frequencyPenalty(0.0)
                .presencePenalty(0.0)
                .temperature(0.00)
                .maxTokens(maxTokens)
                .model("code-davinci-002")
                .echo(true)
                .build();

        List<CompletionChoice> choices = service.createCompletion(completionRequest).getChoices();
        return choices.stream().map(CompletionChoice::getText).collect(Collectors.toList());
    }

    public List<String> calculateTimeComplexity(ChatModel chatModel){
        String context = "%s\n" +
                "\"\"\"\n" +
                "The time complexity of this function is";
        context = String.format(context,chatModel.getInput());
        Duration timeout =  Duration.ofSeconds(60);
        OpenAiService service = new OpenAiService(apiKey,timeout);
        Integer maxTokens = 100;
        List<String> stop = new ArrayList<>();
        stop.add("\\n");
        CompletionRequest completionRequest = CompletionRequest.builder()
                .prompt(context)
                .user("user123")
                .stop(stop)
                .topP(1.0)
                .frequencyPenalty(0.0)
                .presencePenalty(0.0)
                .temperature(0.00)
                .maxTokens(maxTokens)
                .model("text-davinci-003")
                .echo(true)
                .build();

        List<CompletionChoice> choices = service.createCompletion(completionRequest).getChoices();
        return choices.stream().map(CompletionChoice::getText).collect(Collectors.toList());
    }

    public List<String> bugFixer(ChatModel chatModel){

        String context = "##### Fix bugs in the below function\n" +
                "### Buggy %s\n" +
                "%s\n" +
                "### Fixed %s";
        context = String.format(context,chatModel.getLanguage(),chatModel.getInput(),chatModel.getLanguage());
        Duration timeout =  Duration.ofSeconds(120);
        OpenAiService service = new OpenAiService(apiKey,timeout);
        Integer maxTokens = 200;
        List<String> stop = new ArrayList<>();
        stop.add("###");
        CompletionRequest completionRequest = CompletionRequest.builder()
                .prompt(context)
                .user("user123")
                .stop(stop)
                .topP(1.0)
                .frequencyPenalty(0.0)
                .presencePenalty(0.0)
                .temperature(0.00)
                .maxTokens(maxTokens)
                .model("code-davinci-002")
                .echo(true)
                .build();
        List<CompletionChoice> choices = service.createCompletion(completionRequest).getChoices();
        return choices.stream().map(CompletionChoice::getText).collect(Collectors.toList());
    }

    public List<String> explainCode(ChatModel chatModel){

        String context = "%s\n" +
                "\"\"\"\n" +
                "Here's what the above class is doing:";
        context = String.format(context,chatModel.getInput());
        Duration timeout =  Duration.ofSeconds(120);
        OpenAiService service = new OpenAiService(apiKey,timeout);
        Integer maxTokens = 200;
        List<String> stop = new ArrayList<>();
        stop.add("###");
        CompletionRequest completionRequest = CompletionRequest.builder()
                .prompt(context)
                .user("user123")
                .stop(stop)
                .topP(1.0)
                .frequencyPenalty(0.0)
                .presencePenalty(0.0)
                .temperature(0.00)
                .maxTokens(maxTokens)
                .model("code-davinci-002")
                .echo(true)
                .build();
        List<CompletionChoice> choices = service.createCompletion(completionRequest).getChoices();
        return choices.stream().map(CompletionChoice::getText).collect(Collectors.toList());
    }


    public CompletionResult completion(String input, List<String> stop, Integer maxTokens, String model) {


        Duration timeout = Duration.ofSeconds(60);

        OpenAiService service = new OpenAiService(apiKey, timeout);

        CompletionRequest completionRequest = CompletionRequest.builder()
                .prompt(input)
                .user("user123")
                .stop(stop)
                .maxTokens(maxTokens)
                .model(model)
                .echo(false)
                .build();


        String rp = HttpUtil.createPost("https://pms.descartes.digital/openai/v1/completions")
                .timeout(1000*60).body(new JSONObject(completionRequest).toString(), "application/json")
                .bearerAuth(apiKey).execute().body();

        GPTCompletionResult result = new JSONObject(rp).toBean(GPTCompletionResult.class);
        handelResultToken(result);


        return result;
    }

    @Autowired
    private IAiImgCreateRecordService iAiImgCreateRecordService;

    private final Integer openAiImgPrice = 150; //openAi图片单价
    public AiImgCreateRecord simpleCreateImg(String input) {

        int num = 3;
        LoginUser sysUser = (LoginUser) SecurityUtils.getSubject().getPrincipal();

        sysUserService.updateUserScore(sysUser.getId(),openAiImgPrice*-1*3);

        List<String> imgs = this.imgCreate(input, num, "256x256");
        AiImgCreateRecord record = new AiImgCreateRecord();

        SysUser user = sysUserService.getById(sysUser.getId());
        record.setImgs(new JSONArray(imgs).toString());
        record.setIsPush(true);
        record.setPrompt(input);
        record.setUserId(sysUser.getId());
        record.setAvatar(user.getAvatar());
        record.setNickName(user.getNickName());
        record.setNum(num);
        iAiImgCreateRecordService.save(record);

        this.localImgs(record);

        return record;
    }


    @Value(value = "${jeecg.path.upload}")
    private String uploadpath;
    @Value(value = "${open-ai.img-server}")
    private String server;
    @Async
    public void localImgs(AiImgCreateRecord record){

        String ctxPath = uploadpath;
        String fileName = null;
        File file = new File(ctxPath + File.separator + "localImg" + File.separator );
        if (!file.exists()) {
            // 创建文件根目录
            file.mkdirs();
        }



        JSONArray imgs = new JSONArray(record.getImgs());
        JSONArray updateImg = new JSONArray();
        for(int i=0;i<imgs.size();i++){

            String urlPath = imgs.getStr(i);
            // 本地文件夹路径
//            String substring = urlPath.substring(urlPath.lastIndexOf(".")+1);
            // 从服务器端获得图片，并保存到本地
            fileName = FileUtils.saveImageToDisk(getInputStream(urlPath), file.getPath(),"png");
            updateImg.add(server+"localImg/"+fileName);
        }

        record.setImgs(updateImg.toString());

        iAiImgCreateRecordService.updateById(record);

    }


    public static InputStream getInputStream(String urlPath) {
        InputStream inputStream = null;
        HttpURLConnection httpURLConnection = null;
        try {
            URL url = new URL(urlPath);
            httpURLConnection = (HttpURLConnection) url.openConnection();
            // 设置网络连接超时时间
            httpURLConnection.setConnectTimeout(3000);
            // 设置应用程序要从网络连接读取数据
            httpURLConnection.setDoInput(true);
            httpURLConnection.setRequestMethod("GET");
            int responseCode = httpURLConnection.getResponseCode();
            if (responseCode == 200) {
                // 从服务器返回一个输入流
                inputStream = httpURLConnection.getInputStream();
            }
        } catch (MalformedURLException e) {
            e.printStackTrace();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return inputStream;

    }


    public List<String> editImg(EditImg editImg) {


        CreateImageEditRequest createImageEditRequest = CreateImageEditRequest.builder()
                .size("256x256")
                .responseFormat("url")
                .n(3)
                .prompt(editImg.getInput()).build();

        Duration timeout = Duration.ofSeconds(60);

        OpenAiService service = new OpenAiService(apiKey, timeout);

        File f = new File(editImg.getImg1());

        return service.createImageEdit(createImageEditRequest, editImg.getImg1(), editImg.getImg2()).getData().stream().map(i -> i.getUrl()).collect(Collectors.toList());
    }

    public List<String> imgCreate(String input, int n, String size) {

        Duration timeout = Duration.ofSeconds(60);

        OpenAiService service = new OpenAiService(apiKey, timeout);

        CreateImageRequest request = CreateImageRequest.builder()
                .n(n)
                .prompt(input).size(size).build();


        String rp = HttpUtil.createPost("https://pms.descartes.digital/openai/v1/images/generations")
                .timeout(1000*60).body(new JSONObject(request).toString(), "application/json")
                .bearerAuth(apiKey).execute().body();

        ImageResult result = new JSONObject(rp).toBean(ImageResult.class);

        return result.getData().stream().map(i -> i.getUrl()).collect(Collectors.toList());
    }

    public static void test() {

        Duration timeout = Duration.ofSeconds(60);
        OpenAiService service = new OpenAiService(apiKey, timeout);
        List<String> stop = new ArrayList<>();
        stop.add("#");
        stop.add(";");
        CompletionRequest completionRequest = CompletionRequest.builder()
                .prompt("### Postgres SQL tables, with their properties:\n" +
                        "#\n" +
                        "# Employee(id, name, department_id)\n" +
                        "# Department(id, name, address)\n" +
                        "# Salary_Payments(id, employee_id, amount, date)\n" +
                        "#\n" +
                        "### A query to list the names of the departments which employed more than 10 employees in the last 3 months\n" +
                        "SELECT")
//                .user("user123")
                .stop(stop)
                .topP(1.0)
                .frequencyPenalty(0.0)
                .presencePenalty(0.0)
                .maxTokens(300)
                .model("code-davinci-002")
                .echo(false)
                .build();
        service.createCompletion(completionRequest).getChoices().forEach(System.out::println);
    }


    private static int promptLength(String value) {
        int length = 0;
        String chinese = "[\u4e00-\u9fa5]";
        for (int i = 0; i < value.length(); i++) {
            String temp = value.substring(i, i + 1);
            if (temp.matches(chinese)) {
                length += 2;
            } else {
                length += 1;
            }
        }
        return length;
    }

    public static void main(String[] args) {

        String input = "我:你:好xcxvcv:123\n" +
                "AI:12312\n" +
                "我:你好122\n" +
                "\n" +
                "AI:xxcvcv\n";


        String regx = "(我:.*)|(AI:.*)";

        Pattern pat = Pattern.compile(regx);
        Matcher mat = pat.matcher(input);
        List<GPTMessage> list = new ArrayList<>();

        while (mat.find()) {
            String[] str = mat.group().split(":+?");
            GPTMessage message = new GPTMessage();
            message.setRole(str[0]);
            String content = "";
            for (int i = 1; i < str.length; i++) {
                content += str[i];
            }
            message.setContent(content);
            list.add(message);
        }
    }

}
